package com.fanshuai.tree.btree;

import com.fanshuai.tree.bst.BST;

import java.util.ArrayList;
import java.util.List;

/**
 * B+树  B+树为多路搜索树，M表示单个节点最多保存关键字数量
 * 每个节点均有相同的关键字和指针个数
 * 叶节点为数据节点，非叶节点为索引节点。
 * 非叶节点的子树指针p[i]，指向关键字属于[k[i],k[i+1])的子树
 * 叶节点深度相同，所有数据都在叶节点
 * @param <K>
 * @param <V>
 */
public class BPlusTree<K extends Comparable<K>, V> implements BST<K, V> {
    //节点元素数量
    private final static int M = 4;

    class Entry<K, V> {
        public K key;
        public V value;
        public Node next;

        public Entry(K key, V value, Node next) {
            this.key = key;
            this.value = value;
            this.next = next;
        }
    }

    class Node {
        //子树指针
        public Entry<K, V>[] children;
        //元素数量
        public int n;

        public Node(int n) {
            this.children = new Entry[M];
            this.n = n;
        }
    }

    //根节点
    private Node root;
    //元素总数量
    private int size;
    //B+树高度
    private int height;

    public BPlusTree() {
        root = new Node(0);
        size = 0;
        height = 0;
    }

    /**
     * 节点元素满后，分裂节点。将当前节点拆分为2个节点，每个节点元素数量为M/2。2个节点保证元素有序
     * @param node
     * @return
     */
    private Node splitNode(Node node) {
        //拆分节点元素数量
        int n = M - M/2;

        Node next = new Node(n);
        for (int i = 0; i < n; i++) {
            next.children[i] = node.children[M/2 + i];
            node.children[M/2 + i] = null;
        }

        node.n = M/2;

        return next;
    }

    @Override
    public void put(K key, V value) {
        Node nextNode = insert(root, key, value, height);
        if (null == nextNode) {
            return;
        }

        Node node = new Node(2);
        //索引节点
        node.children[0] = new Entry<>(root.children[0].key, null, root);
        node.children[1] = new Entry<>(nextNode.children[0].key, null, nextNode);

        root = node;
        height++;
    }

    /**
     * 递归插入
     * @param node  插入节点
     * @param key
     * @param value
     * @param height  高度  高度为0时，为叶节点
     * @return    节点分裂后的新节点
     */
    private Node insert(Node node, K key, V value, int height) {
        Entry<K, V> entry = null;
        int pos = 0;
        boolean equals = false;
        if (height == 0) { //叶节点  顺序查找
            entry = new Entry<>(key, value, null);

            int i;
            for (i = 0; i < node.n; i++) {
                if (key.equals(node.children[i].key)) {
                    node.children[pos].value = value;  //key相同时，重置value值
                    return null;
                } else if (key.compareTo(node.children[i].key) < 0) {
                    pos = i;
                    break;
                }
            }
            //大于所有元素，放末尾
            if (i == node.n) {
                pos = node.n;
            }
        } else { //非叶节点
            for (int i = 0; i < node.n; i++) {
                if (i == node.n - 1 || key.compareTo(node.children[i + 1].key) < 0) {
                    pos = i;

                    //递归插入
                    Node u = insert(node.children[pos].next, key, value, height - 1);
                    if (null == u) {
                        return null;
                    }

                    //u非空时，返回分裂节点。entry指向分裂节点
                    entry = new Entry<>(u.children[0].key, null, u);
                    //新分裂的节点索引，插入当前索引之后
                    pos++;
                    break;
                }
            }
        }

        if (height == 0) {
            size++;
        }

        for (int j = node.n; j > pos; j--) {
            node.children[j] = node.children[j - 1];
        }
        node.children[pos] = entry;
        node.n++;

        if (node.n == M) {
            //节点已满，分裂节点，并返回新节点
            return splitNode(node);
        } else {
            if (null == entry) {
                return null;
            }
            if (pos == 0) {
                Node node1 = new Node(1);
                node1.children[0] = new Entry<>(entry.key, null, node);
                return node1;
            } else {
                return null;
            }
        }
    }

    /**
     * 查询  从根节点遍历，直到叶节点，在叶节点中查询，返回节点
     * @param key
     * @return
     */
    @Override
    public V get(K key) {
        return get(root, key, height);
    }

    /**
     * 叶节点遍历：顺序查找即可
     * 索引节点遍历：key小于children[i+1].key时，进入第i个节点  遍历到当前节点最后一个元素，进入第n-1个节点
     * @param node
     * @param key
     * @param height
     * @return
     */
    private V get(Node node, K key, int height) {
        if (null == node) {
            return null;
        }

        if (height == 0) {//叶节点
            for (int i = 0; i < node.n; i++) {
                Entry<K, V> entry = node.children[i];
                if (entry.key.equals(key)) {
                    return entry.value;
                }
            }
            return null;
        } else { //非叶节点  寻找小于k[i+1]的节点，进入第i个子节点
            for (int i = 0; i < node.n; i++) {
                if (i == node.n - 1 || key.compareTo(node.children[i + 1].key) < 0) {//尾结点，或者小于i+1
                    return get(node.children[i].next, key, height - 1);
                }
            }
            return null;
        }
    }

    private void deleteEntry(Node node, int pos) {
        for (int j = pos; j < root.n - 1; j++) {
            root.children[j] = root.children[j + 1];
        }
        root.children[root.n - 1] = null;
        root.n--;
    }

    /**
     * 删除  叶节点删除数据。如果节点的首个元素被删除，则上级索引节点对应位置的元素被删除
     * @param key
     */
    @Override
    public void delete(K key) {
        delete(root, key, height);
    }

    /**
     * 递归删除，返回当前节点被删除的索引
     * @param node
     * @param key
     * @return
     */
    private int delete(Node node, K key, int height) {
        if (height == 0) { //数据节点处理
            int pos = -1;
            for (int i = 0; i < node.n; i++) {
                if (key.equals(node.children[i].key)) {
                    pos = i;
                    break;
                }
            }
            if (pos >= 0) {
                deleteEntry(node, pos);
            }

            return pos;
        } else { //索引节点处理
            int pos = -1;
            for (int i = 0; i < node.n; i++) {
                if (i == node.n - 1 || key.compareTo(node.children[i + 1].key) < 0) {
                    pos = i;
                    int index = delete(node.children[i].next, key, height - 1);
                    //子节点首元素被删除，
                    if (index == 0) {
                        deleteEntry(node, pos);
                    }
                    break;
                }
            }
            return pos;
        }
    }

    @Override
    public List<K> keys() {
        List<K> keys = new ArrayList<>(size);
        fillKeys(root, height, keys);
        return keys;
    }

    //B+树有序，中序遍历即可
    private void fillKeys(Node node, int height, List<K> keys) {
        if (height == 0) {
            for (int i = 0; i < node.n; i++) {
                keys.add(node.children[i].key);
            }
        } else {
            for (int i = 0; i < node.n; i++) {
                fillKeys(node.children[i].next, height - 1, keys);
            }
        }
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean empty() {
        return size == 0;
    }
}
